#pragma once
#include <iostream>

namespace my_rbt
{
	// 红黑树 我们只掌握插入，删除暂且不做考虑
	enum COLOR
	{
		RED,
		BLACK
	};
	template <typename K, typename V>
	struct RBNode
	{
		// 默认的颜色为什么是红色而不是黑色？
		// 默认颜色是红色在调整时，代价相对小一些
		// 默认是黑色，调整的代价比较大
		RBNode(K key, V val, COLOR col = RED)
			: _left(nullptr),
			_right(nullptr),
			_parent(nullptr),
			_col(col),
			_key(key),
			_val(val)
		{
		}
		RBNode<K, V>* _left;
		RBNode<K, V>* _right;
		RBNode<K, V>* _parent;
		COLOR _col;
		K _key;
		V _val;
	};
	template <typename K, typename V>
	class RBTree
	{
	private:
		typedef RBNode<K, V> Node;
		typedef RBTree<K, V> Self;

		Node* get_node(K key, V val, COLOR col)
		{
			return new Node(key, val, col);
		}
		void LRotate(Node* cur)
		{
			Node* parent = cur->_parent;
			Node* curL = cur->_left;
			Node* pparent = parent->_parent;

			parent->_right = curL;
			if (curL)
				curL->_parent = parent;

			cur->_left = parent;
			parent->_parent = cur;

			if (parent == _root)
			{
				_root = cur;
				cur->_parent = nullptr;
			}
			else
			{
				if (pparent->_left == parent)
					pparent->_left = cur;
				else
					pparent->_right = cur;
				cur->_parent = pparent;
			}
		}
		void RRotate(Node* cur)
		{
			Node* parent = cur->_parent;
			Node* curR = cur->_right;
			Node* pparent = parent->_parent;

			parent->_left = curR;
			if (curR)
				curR->_parent = parent;

			cur->_right = parent;
			parent->_parent = cur;

			if (parent == _root)
			{
				_root = cur;
				_root->_parent = nullptr;
			}
			else
			{
				if (pparent->_left == parent)
					pparent->_left = cur;
				else
					pparent->_right = cur;
				cur->_parent = pparent;
			}
		}

		int GetBlackL(Node* root)
		{
			// 这么写是获取子树上的黑色节点，而非支路上的黑色节点
			//if (root == nullptr)
			//	return 1;
			//int left = GetBlack(root->_left);
			//int right = GetBlack(root->_right);
			//int cnt = 0;
			//if (root->_col == BLACK)
			//	cnt = 1;
			//return left + right + cnt;
			int cnt = 0;
			while (root)
			{
				if (root->_col == BLACK)
					cnt++;
				root = root->_left;
			}
			return cnt;

		}
		int GetBlackR(Node* root)
		{
			int cnt = 0;
			while (root)
			{
				if (root->_col == BLACK)
					cnt++;
				root = root->_right;
			}
			return cnt;
		}
		// base：最左侧分支上 黑色节点的数目
		// pre：判断是否是二叉搜索树
		// cnt: 当前支路上的黑色节点数
		bool _IsRBTree(Node* root, int cnt, int& base, Node*& pre)
		{
			if (root == nullptr)
			{
				if (base == 0)
					cnt = base;
				else
				{
					if (cnt != base)
					{
						std::cout << "分支上黑色节点不同" << std::endl;
						return false;
					}
				}
				return true;
			}
			if (root->_col == BLACK)
				cnt++;

			if (root->_col == RED && root->_left && root->_right && !(root->_left->_col == BLACK && root->_right->_col == BLACK))
			{
				std::cout << "!(root->_left->_col == BLACK && root->_right->_col == BLACK)" << std::endl;
				std::cout << "root->_col = " << root->_col << std::endl;
				std::cout << "root->_left->_col = " << root->_left->_col;
				std::cout << " root->_right->_col = " << root->_right->_col << std::endl;
				return false;
			}


			bool left = _IsRBTree(root->_left, cnt, base, pre);
			if (pre == nullptr)
				pre = root;
			else
			{
				if (root->_key <= pre->_key)
				{
					std::cout << "不是二叉搜索树" << std::endl;
					return false;
				}
			}

			bool right = _IsRBTree(root->_right, cnt, base, pre);

			return left && right;
		}

	public:
		bool IsRBTree()
		{
			if (_root->_col != BLACK)
			{
				std::cout << "_root->_col != BLACK" << std::endl;
				return false;
			}
			int base = 0;
			Node* pre = nullptr;
			return _IsRBTree(_root, 0, base, pre);
		}

		bool insert(K key, V val)
		{
			if (_root == nullptr)
			{
				_root = get_node(key, val, BLACK);
				return true;
			}
			Node* cur = _root;
			Node* parent = cur;
			while (cur)
			{
				if (key > cur->_key)
				{
					parent = cur;
					cur = cur->_right;
				}
				else if (key < cur->_key)
				{
					parent = cur;
					cur = cur->_left;
				}
				else
					return false;
			}
			cur = get_node(key, val, RED);
			if (key > parent->_key)
				parent->_right = cur;
			else
				parent->_left = cur;
			cur->_parent = parent;
			// 向上调节 连续红色才调整
			while (parent && parent->_col == RED)
			{
				Node* pparent = parent->_parent;
				Node* uncle = nullptr;
				if (pparent)
				{
					if (pparent->_left == parent)
						uncle = pparent->_right;
					else
						uncle = pparent->_left;
				}
				if (cur->_col == RED && parent->_col == RED)
				{
					if (uncle && uncle->_col == RED /*&& pparent->_col == BLACK*/)
					{
						parent->_col = BLACK;
						uncle->_col = BLACK;
						pparent->_col = RED;


						cur = pparent;
						parent = cur->_parent;
					}
					else
					{
						//     pp
						//   p
						// c
						//     p 
						//  c     pp
						if (pparent->_left == parent && parent->_left == cur)
						{
							RRotate(parent);
							cur->_col = pparent->_col = RED;
							parent->_col = BLACK;
							break;
						}
						else if (pparent->_right == parent && parent->_right == cur)
						{
							LRotate(parent);
							cur->_col = pparent->_col = RED;
							parent->_col = BLACK;
							break;

						}
						else if (pparent->_left == parent && parent->_right == cur)
						{
							LRotate(cur);
							RRotate(cur);
							pparent->_col = RED;
							parent->_col = RED;
							cur->_col = BLACK;

							break;

						}
						else if (pparent->_right == parent && parent->_left == cur)
						{
							RRotate(cur);
							LRotate(cur);
							cur->_col = BLACK;
							pparent->_col = RED;
							parent->_col = RED;
							break;

						}
					}
				}

			}
			_root->_col = BLACK;
			return true;
		}

        Node* find(const K& key)
        {
            Node* cur = _root;
            while(cur)
            {
                if(key > cur->_key)
                    cur = cur->_right;
                else if (key < cur->_key)
                    cur = cur->_left;
                else
                    return cur;
            }
            return nullptr;
        }

	private:
		Node* _root = nullptr;
	};
}